import logging
from pathlib import PurePath
from typing import Callable, Collection, Set, Type

from common import OperatingSystem
from common.agent_events import AgentEventTag
from common.credentials import Credentials
from common.tags import (
    BRUTE_FORCE_T1110_TAG,
    COMMAND_AND_SCRIPTING_INTERPRETER_T1059_TAG,
    FILE_AND_DIRECTORY_PERMISSIONS_MODIFICATION_T1222_TAG,
    INGRESS_TOOL_TRANSFER_T1105_TAG,
    REMOTE_SERVICES_T1021_TAG,
)
from common.types import NetworkPort
from infection_monkey.exploit.tools import (
    IRemoteAccessClient,
    RemoteAuthenticationError,
    RemoteCommandExecutionError,
    RemoteFileCopyError,
)
from infection_monkey.i_puppet import TargetHost

from .ssh_client import SSHClient
from .ssh_options import SSHOptions

logger = logging.getLogger(__name__)
LOGIN_TAGS = {BRUTE_FORCE_T1110_TAG, REMOTE_SERVICES_T1021_TAG}
COPY_FILE_TAGS = {
    FILE_AND_DIRECTORY_PERMISSIONS_MODIFICATION_T1222_TAG,
    INGRESS_TOOL_TRANSFER_T1105_TAG,
    REMOTE_SERVICES_T1021_TAG,
}
EXECUTION_TAGS = {COMMAND_AND_SCRIPTING_INTERPRETER_T1059_TAG, REMOTE_SERVICES_T1021_TAG}

SSH_PORTS = [NetworkPort(22)]


class SSHRemoteAccessClient(IRemoteAccessClient):
    def __init__(
        self,
        host: TargetHost,
        options: SSHOptions,
        build_command: Callable[[PurePath], str],
        ssh_client: SSHClient,
    ):
        self._host = host
        self._options = options
        self._build_command = build_command
        self._ssh_client = ssh_client

    def login(self, credentials: Credentials, tags: Set[AgentEventTag]):
        tags.update(LOGIN_TAGS)

        try:
            for port in SSH_PORTS:
                self._ssh_client.connect(
                    self._host, credentials, port, timeout=self._options.connect_timeout
                )
        except Exception as err:
            error_message = f"Failed to authenticate over SSH with {credentials}: {err}"
            raise RemoteAuthenticationError(error_message)

    def _raise_if_not_authenticated(self, error_type: Type[Exception]):
        if not self._ssh_client.connected():
            raise error_type(
                "This operation cannot be performed until authentication is successful"
            )

    def get_os(self) -> OperatingSystem:
        return OperatingSystem.LINUX

    def execute_agent(self, agent_binary_path: PurePath, tags: Set[AgentEventTag]):
        self._raise_if_not_authenticated(RemoteCommandExecutionError)

        tags.update(EXECUTION_TAGS)

        try:
            self._ssh_client.execute_command(
                self._build_command(remote_agent_binary_destination_path=agent_binary_path)
            )
        except Exception as err:
            raise RemoteCommandExecutionError(err)

    def copy_file(self, file: bytes, destination_path: PurePath, tags: Set[AgentEventTag]):
        self._raise_if_not_authenticated(RemoteFileCopyError)

        logger.debug(
            f"Trying to copy monkey file to [{destination_path}] on victim {self._host.ip}"
        )

        logger.info(f"Attempting to copy the monkey agent binary to {self._host.ip}")
        tags.update(COPY_FILE_TAGS)

        try:
            self._ssh_client.copy_file(file, destination_path)
        except Exception as err:
            raise RemoteFileCopyError(f"Failed to copy the agent binary to the victim: {err}")

    def get_writable_paths(self) -> Collection[PurePath]:
        return []
